import pandas as pd

from StressAna.Lib.Model import Globalvariables

columns = ['sta', 'length', 'wallThick', 'holeDepth', 'initialPress']

'''
获取台站属性df
'''


def sIDAttrInit():
    if Globalvariables.LABELS_DICT.empty:
        data_list = []
        for index, sta in enumerate(Globalvariables.STATION_dict):
            data = {}
            data['sta'] = sta
            data['length'] = Globalvariables.STATION_CD[index]
            data['wallThick'] = Globalvariables.STATION_BH[index]
            data['holeDepth'] = Globalvariables.STATION_KS[index]
            data['initialPress'] = Globalvariables.STATION_CSY[index]
            data_list.append(data)
        Globalvariables.LABELS_DICT = pd.DataFrame(data=data_list, columns=columns)
        # df.to_json('Labels.json', orient='index')


'''
获取单个台站属性
'''


def get_sIDAttr(sID):
    iAttrList = []
    df = Globalvariables.LABELS_DICT
    sta_list = df['sta'].tolist()
    if sID in sta_list:
        length = df[(df['sta'] == sID)]['length'].values[0]
        iAttrList.append(length)
        wallThick = df[(df['sta'] == sID)]['wallThick'].values[0]
        iAttrList.append(wallThick)
        holeDepth = df[(df['sta'] == sID)]['holeDepth'].values[0]
        iAttrList.append(holeDepth)
        initialPress = df[(df['sta'] == sID)]['initialPress'].values[0]
        iAttrList.append(initialPress)
    return ','.join(iAttrList)


def labelsInit(sIDs):
    vs = ['CD', 'BH', 'KS', 'CSY']
    labels = []
    sIDAttrInit()
    df = Globalvariables.LABELS_DICT
    for sID in sIDs:
        if sID in df['sta'].to_list():
            iLabelStrList = [sID]
            if 'CD' in vs:
                length = df[(df['sta'] == sID)]['length'].values[0]
                iLabelStrList.append(length)
            if 'BH' in vs:
                wallThick = df[(df['sta'] == sID)]['wallThick'].values[0]
                iLabelStrList.append(wallThick)
            if 'KS' in vs:
                holeDepth = df[(df['sta'] == sID)]['holeDepth'].values[0]
                iLabelStrList.append(holeDepth)
            if 'CSY' in vs:
                initialPress = df[(df['sta'] == sID)]['initialPress'].values[0]
                iLabelStrList.append(initialPress)
            labels.append(','.join(iLabelStrList))
        else:
            labels.append(sID)
    # print(labels)
    return labels


if __name__ == '__main__':
    labelsInit(['G12-1', '51-2'])
